import random
from common_import import *
import model_p4
import cal_range


# 定义目标函数 f
def objective_function(x):
    X = np.array(x).reshape(1, -1)
    result = model_p4.combined_prediction(X)
    return result[0]  # 示例目标函数，可以根据需求更改


# 初始化种群
def initialize_population(
    pop_size=50,
    dimensions=295,
):
    population = []
    length = 0
    if pop_size == 50 and dimensions == 295:
        molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
        selected_features = molecular_descriptor.loc[
            constants.indice_50[:length], constants.feature_295
        ]
        selected_features_list = selected_features.values.tolist()
        # return selected_features_list
        population.extend(selected_features_list)
    for _ in range(50 - length):
        individual = [
            random.uniform(cal_range.min_values[i], cal_range.max_values[i])
            for i in range(dimensions)
        ]
        population.append(individual)
    return population


# 评估种群
def evaluate_population(population):
    X = np.array(population)
    return model_p4.combined_prediction(X)
    # return [objective_function(individual) for individual in population]


# 锦标赛选择
def tournament_selection(population, fitness, num_parents, tournament_size=3):
    parents = []
    for _ in range(num_parents):
        selected = random.sample(range(len(population)), tournament_size)
        best = max(selected, key=lambda idx: fitness[idx])
        parents.append(population[best])
    return parents


# 多点交叉操作
def crossover(parents, offspring_size):
    offspring = []
    for k in range(offspring_size):
        parent1_idx = k % len(parents)
        parent2_idx = (k + 1) % len(parents)
        crossover_points = sorted(random.sample(range(1, len(parents[0])), 2))
        offspring.append(
            parents[parent1_idx][: crossover_points[0]]
            + parents[parent2_idx][crossover_points[0] : crossover_points[1]]
            + parents[parent1_idx][crossover_points[1] :]
        )
    return offspring


# 自适应变异操作
def adaptive_mutation(
    offspring, mutation_rate, lower_bound, upper_bound, current_gen, max_gen
):
    adaptive_rate = mutation_rate * (1 - current_gen / max_gen)  # 自适应变异率
    for individual in offspring:
        if random.random() < adaptive_rate:  # 检查是否进行变异
            # 随机选择变异的位置
            mutation_index = random.randint(0, len(individual) - 1)
            # 在指定范围内对该位置进行变异
            individual[mutation_index] = random.uniform(
                cal_range.min_values[mutation_index],
                cal_range.max_values[mutation_index],
            )
    return offspring


# 遗传算法主循环
def genetic_algorithm(
    pop_size,
    dimensions,
    num_generations,
    num_parents,
    mutation_rate,
    lower_bound,
    upper_bound,
):
    # 初始化种群
    population = initialize_population(pop_size, dimensions)

    # 记录每一代的最佳适应度和平均适应度
    best_fitness_per_generation = []
    average_fitness_per_generation = []

    for generation in range(num_generations):
        # 评估种群
        fitness = evaluate_population(population)
        best_fitness = max(fitness)  # 找到当前种群的最佳适应度
        average_fitness = sum(fitness) / len(fitness)  # 计算当前种群的平均适应度

        # 记录每一代的最佳适应度和平均适应度
        best_fitness_per_generation.append(best_fitness)
        average_fitness_per_generation.append(average_fitness)

        if (generation - 1) % 100 == 0:
            print(
                f"Generation {generation}: Best fitness = {best_fitness}, Average fitness = {average_fitness}"
            )

        # 精英保留
        elite_size = 2
        elite_indices = sorted(
            range(len(fitness)), key=lambda i: fitness[i], reverse=True
        )[:elite_size]
        elites = [population[i] for i in elite_indices]

        # 选择父母
        parents = tournament_selection(population, fitness, num_parents)

        # 生成后代
        offspring_size = pop_size - len(elites)
        offspring = crossover(parents, offspring_size)

        # 自适应变异操作
        offspring = adaptive_mutation(
            offspring,
            mutation_rate,
            lower_bound,
            upper_bound,
            generation,
            num_generations,
        )

        # 创建新的种群
        population = elites + offspring

    # 绘制遗传算法的变化曲线图

    # 创建一个包含两个子图的图形，并设置大小
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))

    # 绘制第一个子图：最佳适应度
    ax1.plot(
        range(num_generations),
        best_fitness_per_generation,
        linestyle="-",
        color="b",
        label="最佳适应度",
    )
    ax1.set_title("每代最佳适应度", fontsize=14)
    ax1.set_xlabel("代数", fontsize=12)
    ax1.set_ylabel("适应度", fontsize=12)
    ax1.legend()
    ax1.grid(True)

    # 绘制第二个子图：平均适应度
    ax2.plot(
        range(num_generations),
        average_fitness_per_generation,
        linestyle="-",
        color="r",
        label="平均适应度",
    )
    ax2.set_title("每代平均适应度", fontsize=14)
    ax2.set_xlabel("代数", fontsize=12)
    ax2.set_ylabel("适应度", fontsize=12)
    ax2.legend()
    ax2.grid(True)

    # 调整子图之间的间距
    plt.tight_layout()

    # 显示图形
    plt.show()


# 最佳解
# best_solution = max(population, key=objective_function)
# print("Best solution found:", best_solution)
# print("Best solution fitness:", objective_function(best_solution))


if __name__ == "__main__":
    # 参数设置
    pop_size = 50
    dimensions = 295
    num_generations = 20000
    num_parents = 20
    mutation_rate = 0.2
    lower_bound, upper_bound = -10.0, 10.0

    # 运行遗传算法
    genetic_algorithm(
        pop_size,
        dimensions,
        num_generations,
        num_parents,
        mutation_rate,
        lower_bound,
        upper_bound,
    )
